import yaml
import re


def load_config():
    """加载配置文件"""
    with open('config/run_config.yaml', 'r', encoding='utf-8') as f:
        run_config = yaml.safe_load(f)
    with open('config/thcp_network.yaml', 'r', encoding='utf-8') as f:
        network_config = yaml.safe_load(f)
    return run_config['app'], network_config['thcp']


def parse_nodelist(nodelist_str):
    """解析nodelist的多种格式，返回节点列表"""
    nodes = []
    
    # 去除空格
    nodelist_str = nodelist_str.strip()
    
    # 格式1: "cn[6831-6836,7034,7111]" 或 格式3: "cn[6831-6836],cn[7034,7111]"
    if 'cn[' in nodelist_str:
        # 找到所有cn[...]形式的片段
        pattern = r'cn\[([^\]]+)\]'
        matches = re.findall(pattern, nodelist_str)
        
        for match in matches:
            parts = match.split(',')
            for part in parts:
                part = part.strip()
                if '-' in part:
                    start, end = part.split('-')
                    for i in range(int(start), int(end) + 1):
                        nodes.append('cn' + str(i))
                else:
                    nodes.append('cn' + part)
    
    # 格式2: "cn6831 cn6832 cn6833"
    elif ' ' in nodelist_str:
        nodes = nodelist_str.split()
    
    else:
        # 单个节点
        nodes = [nodelist_str]
    
    return nodes


def get_which_node(proc_num, node_list, proc_per_node, distribution):
    """根据进程号获取所在节点"""
    if distribution == 'block':
        node_index = proc_num // proc_per_node
    else:  # cycle
        node_index = proc_num % len(node_list)
    
    return node_list[node_index]


def get_which_frame(node_name, node_num_per_frame):
    """根据实际cn节点获取机框号"""
    node_num = int(node_name.replace('cn', ''))
    frame = node_num // node_num_per_frame
    return frame


def get_which_cabinet(node_name, node_num_per_cabinet):
    """根据实际cn节点获取机柜号"""
    node_num = int(node_name.replace('cn', ''))
    cabinet = node_num // node_num_per_cabinet
    return cabinet


def get_network_type(src, dst):
    """根据源和目的进程号获取网络层次类型"""
    run_config, network_config = load_config()
    
    # 解析节点列表
    node_list = parse_nodelist(run_config['nodelist'])
    proc_per_node = run_config['proc_per_node']
    distribution = run_config['distribution']
    
    # 获取网络配置
    node_num_per_frame = network_config['node_num_per_frame']
    node_num_per_cabinet = network_config['node_num_per_cabinet']
    
    # 获取节点号
    node_src = get_which_node(src, node_list, proc_per_node, distribution)
    node_dst = get_which_node(dst, node_list, proc_per_node, distribution)
    
    # 获取机框号
    frame_src = get_which_frame(node_src, node_num_per_frame)
    frame_dst = get_which_frame(node_dst, node_num_per_frame)
    
    # 获取机柜号
    cabinet_src = get_which_cabinet(node_src, node_num_per_cabinet)
    cabinet_dst = get_which_cabinet(node_dst, node_num_per_cabinet)
    
    # 判断网络层次类型
    if node_src == node_dst:
        network_type = 0
    elif frame_src == frame_dst:
        network_type = 1
    elif cabinet_src == cabinet_dst:
        network_type = 2
    else:
        network_type = 3
    
    return network_type


if __name__ == '__main__':
    # 测试
    src = 0
    dst = 1
    
    run_config, network_config = load_config()
    node_list = parse_nodelist(run_config['nodelist'])
    proc_per_node = run_config['proc_per_node']
    distribution = run_config['distribution']
    node_num_per_frame = network_config['node_num_per_frame']
    node_num_per_cabinet = network_config['node_num_per_cabinet']
    
    # 获取src信息
    node_src = get_which_node(src, node_list, proc_per_node, distribution)
    frame_src = get_which_frame(node_src, node_num_per_frame)
    cabinet_src = get_which_cabinet(node_src, node_num_per_cabinet)
    
    # 获取dst信息
    node_dst = get_which_node(dst, node_list, proc_per_node, distribution)
    frame_dst = get_which_frame(node_dst, node_num_per_frame)
    cabinet_dst = get_which_cabinet(node_dst, node_num_per_cabinet)
    
    # 获取网络类型
    result = get_network_type(src, dst)
    
    print(f"src: {src}, node: {node_src}, frame: {frame_src}, cabinet: {cabinet_src}")
    print(f"dst: {dst}, node: {node_dst}, frame: {frame_dst}, cabinet: {cabinet_dst}")
    print(f"network_type: {result}")
